// -----------------------------
// Author: 		Dhaenens Gert	
// Creation: 	10/16/2014 9:40:38 PM				
// -----------------------------

#include "ComputePass.h"

#include "Application.h"
#include "Window.h"
#include "RenderManager.h"
#include "FileManager.h"

#include "Scene.h"
#include "LightActor.h"
#include "Camera.h"

#define DISPATCH_NUM_THREADS_X	16
#define DISPATCH_NUM_THREADS_Y	16

struct LightInfo
{
	Vector3	Position;
	Float	Range;
	Vector3	Direction;
	Float	Intensity;
	UInt32	LightType;
	Vector3	Padding;
};

struct DispatchBuffer
{
	UInt32	LightCount;
	Vector3	CameraPosition;
};

ComputePass::ComputePass(RenderManager* lpRenderManager, const RenderTargets& targets)
	: RenderPass(lpRenderManager, "Compute", targets)
	, m_lpLightShader(nullptr)
	, m_lpDispatchBuffer(nullptr)
	, m_lpLightBuffer(nullptr)
	, m_ThreadCountX(0)
	, m_ThreadCountY(0)
	, m_ThreadCountZ(0)
{
}

ComputePass::~ComputePass()
{
	NullRelease(&m_lpLightBuffer);
	NullRelease(&m_lpDispatchBuffer);
	NullRelease(&m_lpLightShader);
}

void ComputePass::Init(Application* lpApp)
{
	TString lightShader = lpApp->GetFiles()->GetFileName(_T("Shader_LightingPass"));

	m_lpLightShader = ShaderCompiler::LoadComputeFromFile(m_lpRenderManager->GetDevice(), lightShader.c_str());
	m_lpDispatchBuffer = IBuffer::Create(m_lpRenderManager->GetDevice(), BufferType::Constant, Usage::Default, 0, sizeof(DispatchBuffer), 1, nullptr);

	m_ThreadCountX = (lpApp->GetWindow()->GetWidth() + DISPATCH_NUM_THREADS_X) / DISPATCH_NUM_THREADS_X;
	m_ThreadCountY = (lpApp->GetWindow()->GetHeight() + DISPATCH_NUM_THREADS_Y) / DISPATCH_NUM_THREADS_Y;
	m_ThreadCountZ = 1;
}
void ComputePass::Apply(Application* lpApp, Scene* lpScene)
{
	IRenderDevice* lpDevice = m_lpRenderManager->GetDevice();

	// Get the Lights
	vector<LightActor*> lightActors = lpScene->GetActorsWithFlagsAs<LightActor>(Actor::Light);

	// Make sure we actually have lights
	if (lightActors.size() < 1)
		return;

	// Build the Light Buffer Data
	vector<LightInfo> lightInfos;
	for (LightActor* lpActor : lightActors)
	{
		LightInfo lightInfo = {};
		
		lightInfo.Position = lpActor->GetPosition();
		lightInfo.Direction = lpActor->GetForward();
		lightInfo.Range = lpActor->GetRange();
		lightInfo.Intensity = lpActor->GetIntensity();
		lightInfo.LightType = lpActor->GetLightType();

		lightInfos.push_back(lightInfo);
	}

	// Update the Light Buffer
	// If it does not exist yet, or we have a different
	// amount of lights, recreate it
	if (!m_lpLightBuffer || lightInfos.size() != m_lpLightBuffer->GetCount())
	{
		NullRelease(&m_lpLightBuffer);
		m_lpLightBuffer = IBuffer::Create(lpDevice, BufferType::Structured, Usage::Default, 0, sizeof(LightInfo), (UInt32) lightInfos.size(), lightInfos.data());
	}
	else
		m_lpLightBuffer->Update(lightInfos.data(), sizeof(LightInfo)* (UInt32) lightInfos.size());

	// Bind the Light Buffer
	m_lpLightShader->SetBufferAsResource(5, m_lpLightBuffer);

	// Bind the Compute Shader
	lpDevice->SetComputeShader(m_lpLightShader);
	m_lpLightShader->SetBufferAsConstant(0, m_lpDispatchBuffer);

	// Bind the Target Textures
	m_lpLightShader->SetTargetAsResource(0, m_Targets.PositionTarget);
	m_lpLightShader->SetTargetAsResource(1, m_Targets.NormalTarget);
	m_lpLightShader->SetTargetAsResource(2, m_Targets.AlbedoTarget);
	m_lpLightShader->SetTargetAsResource(3, m_Targets.AmbientTarget);
	m_lpLightShader->SetStencilAsResource(4, m_Targets.DepthStencil);

	// Bind the Output Texture
	m_lpLightShader->SetTargetAsUav(0, m_Targets.RawRenderTarget3D);

	// Update the Dispatch Buffer
	DispatchBuffer dispatchBuffer = {};
	dispatchBuffer.LightCount = (UInt32) lightInfos.size();
	dispatchBuffer.CameraPosition = lpScene->GetCamera()->GetPosition();
	m_lpDispatchBuffer->Update(&dispatchBuffer, sizeof(DispatchBuffer));

	// Dispatch the Compute Shader
	lpDevice->Dispatch(m_ThreadCountX, m_ThreadCountY, m_ThreadCountZ);

	// Unbind the Compute Shader
	lpDevice->SetComputeShader(nullptr);

	// Unbind the Targets
	m_lpLightShader->ClearResource(0);
	m_lpLightShader->ClearResource(1);
	m_lpLightShader->ClearResource(2);
	m_lpLightShader->ClearResource(3);
	m_lpLightShader->ClearResource(4);
	m_lpLightShader->ClearResource(5);
	m_lpLightShader->ClearUav(0);
}
void ComputePass::UpdateTargets(Application* lpApp, const RenderTargets& targets)
{
	RenderPass::UpdateTargets(lpApp, targets);

	// Update the Thread Counts
	m_ThreadCountX = (lpApp->GetWindow()->GetWidth() + DISPATCH_NUM_THREADS_X) / DISPATCH_NUM_THREADS_X;
	m_ThreadCountY = (lpApp->GetWindow()->GetHeight() + DISPATCH_NUM_THREADS_Y) / DISPATCH_NUM_THREADS_Y;
	m_ThreadCountZ = 1;
}